from typing import Optional, Any, List

from openai import OpenAI
from llama_index.core.bridge.pydantic import Field, PrivateAttr

from llama_index.core.base.embeddings.base import BaseEmbedding

from app.core import config



class OpenAIEmbedding(BaseEmbedding):
    api_key: str = Field(default='ollama',description="The OpenAI API key.")
    api_base: Optional[str] = Field(
        default='https://api.openai.com/v1', description="The base URL for OpenAI API."
    )
    model: Optional[str] = Field(
        default='text-embedding-v3', description="向量模型."
    )
    _client: Optional[OpenAI] = PrivateAttr()
    dimensions: Optional[int] = Field(
        default=None, description="维度,有些模型不支持设置该参数"
    )

    def __init__(
            self,
            api_base: Optional[str] = "https://api.openai.com/v1",
            api_key: Optional[str] = 'your-openai-api-key-here',
            model: Optional[str] = "text-embedding-v3",
            dimensions: int= 0,
            **kwargs,
    ):
        """"""
        # OpenAIEmbeddings(openai_api_base=api_base,
        #                  openai_api_key=api_key,
        #                  model=model,
        #                  )
        super().__init__(
            api_key=api_key,
            api_base=api_base,
            model=model,
            **kwargs
        )
        self.dimensions = dimensions
        self._client = OpenAI(
            api_key=api_key,  # 如果您没有配置环境变量，请在此处用您的API Key进行替换
            base_url=api_base,  # 百炼服务的base_url
        )

    async def _aget_query_embedding(self, query: str) -> List[float]:
        """"""
        return await  self.aget_embedding(
            text=query,

        )

    def _get_query_embedding(self, query: str) -> List[float]:
        """Get query embedding."""
        return  self.get_embedding(
            text=query,
        )

    def _get_text_embedding(self, text: str) -> List[float]:
        """Get text embedding."""
        return  self.get_embedding(
            text=text,
        )

    async def aget_embedding(
            self, text: str
    ) -> List[float]:
        """Asynchronously get embedding.

        NOTE: Copied from OpenAI's embedding utils:
        https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py

        Copied here to avoid importing unnecessary dependencies
        like matplotlib, plotly, scipy, sklearn.

        """
        text = text.replace("\n", " ")

        res = (await self._client.embeddings.create(
            model=self.model,  # "text-embedding-v3",#	1,024（默认）、768或512
            dimensions=self.dimensions,
            input=[text],
            encoding_format="float"
        ))
        return res.data[0].embedding


    def get_embedding(self, text: str) -> List[float]:
        """Get embedding.

        NOTE: Copied from OpenAI's embedding utils:
        https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py

        Copied here to avoid importing unnecessary dependencies
        like matplotlib, plotly, scipy, sklearn.

        """
        text = text.replace("\n", " ")

        from openai.types import CreateEmbeddingResponse
        res:CreateEmbeddingResponse = (self._client.embeddings.create(
            model=self.model,  # "text-embedding-v3",#	1,024（默认）、768或512
            dimensions=self.dimensions,
            input=[text],
            encoding_format="float"
        ))
        # print(res.data[0].embedding)
        return res.data[0].embedding

